{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "73df3d10-94a0-4807-a7ef-17368c2827f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: scipy in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (1.15.2)\n",
      "Requirement already satisfied: tenacity in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (9.0.0)\n",
      "Requirement already satisfied: tiktoken in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (0.9.0)\n",
      "Requirement already satisfied: termcolor in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (2.5.0)\n",
      "Requirement already satisfied: openai in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (1.61.1)\n",
      "Requirement already satisfied: requests in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (2.32.3)\n",
      "Requirement already satisfied: numpy<2.5,>=1.23.5 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from scipy) (1.26.4)\n",
      "Requirement already satisfied: regex>=2022.1.18 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from tiktoken) (2024.11.6)\n",
      "Requirement already satisfied: anyio<5,>=3.5.0 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from openai) (4.6.2)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from openai) (1.9.0)\n",
      "Requirement already satisfied: httpx<1,>=0.23.0 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from openai) (0.27.0)\n",
      "Requirement already satisfied: jiter<1,>=0.4.0 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from openai) (0.9.0)\n",
      "Requirement already satisfied: pydantic<3,>=1.9.0 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from openai) (2.9.2)\n",
      "Requirement already satisfied: sniffio in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from openai) (1.3.0)\n",
      "Requirement already satisfied: tqdm>4 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from openai) (4.67.1)\n",
      "Requirement already satisfied: typing-extensions<5,>=4.11 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from openai) (4.12.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from requests) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from requests) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from requests) (2.3.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from requests) (2025.1.31)\n",
      "Requirement already satisfied: httpcore==1.* in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from httpx<1,>=0.23.0->openai) (1.0.2)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai) (0.14.0)\n",
      "Requirement already satisfied: annotated-types>=0.6.0 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from pydantic<3,>=1.9.0->openai) (0.7.0)\n",
      "Requirement already satisfied: pydantic-core==2.23.4 in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from pydantic<3,>=1.9.0->openai) (2.23.4)\n",
      "Requirement already satisfied: colorama in d:\\java\\miniconda3\\envs\\langchain\\lib\\site-packages (from tqdm>4->openai) (0.4.6)\n"
     ]
    }
   ],
   "source": [
    "# !pip install scipy tenacity tiktoken termcolor openai requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2f993515-aed5-4a53-b0f4-df41ee17577b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 终端执行\n",
    "# Unix系统\n",
    "# export OPENAI_API_KEY = \"sk-y7DHfp9fzuCxOVm2158638099f9541D3833aB4F4Ed674aCf\"\n",
    "# window系统\n",
    "!set OPENAI_API_KEY = \"sk-y7DHfp9fzuCxOVm2158638099f9541D3833aB4F4Ed674aCf\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "1ad789de-4bb8-4c53-830a-6dc7d385503d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import requests\n",
    "import os\n",
    "import random\n",
    "from tenacity import retry, wait_random_exponential, stop_after_attempt\n",
    "from termcolor import colored\n",
    "\n",
    "GPT_MODEL = \"gpt-4o-mini\"\n",
    "# api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "api_key = \"sk-y7DHfp9fzuCxOVm2158638099f9541D3833aB4F4Ed674aCf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "17346cbe-3ccf-4abe-a5df-a4e629a7ab4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 请求OpenAI接口补全消息\n",
    "# 使用retry库，指定在请求失败时的重试策略\n",
    "# wait_random_exponential：指数等待，时间间隔最大为40s，最多重试3次\n",
    "\n",
    "# 智谱模型 Function Calling\n",
    "# 参考： https://www.bigmodel.cn/dev/howuse/functioncall\n",
    "\n",
    "# deepseek-chat 模型 Function Calling\n",
    "# 注意：当前版本 deepseek-chat 模型 Function Calling 功能效果不稳定，会出现循环调用、空回复的情况。\n",
    "# 参考： https://api-docs.deepseek.com/zh-cn/guides/function_calling\n",
    "\n",
    "### 使用 智谱模型 Function Calling 和 deepseek-chat 模型 Function Calling 注意替换模型名称，调用地址，和API KEY\n",
    "@retry(wait=wait_random_exponential(multiplier=1, max=40), stop=stop_after_attempt(3))\n",
    "def chat_completion_request(messages, functions = None, function_call = None, model = GPT_MODEL):\n",
    "    headers = {\n",
    "        \"Content-Type\": \"application/json\",\n",
    "        \"Authorization\": \"Bearer \" + api_key\n",
    "    }\n",
    "    # 设定请求的JSON数据，包括GPT模型名和要进行补全的消息\n",
    "    json_data = {\"model\": model, \"messages\": messages}\n",
    "\n",
    "    # 如果传入functions，则在消息体中加入配置\n",
    "    if functions is not None:\n",
    "        json_data.update({\"tools\": functions})\n",
    "\n",
    "    # 如果传入了function_call，则在消息体中加入该配置\n",
    "    if function_call is not None:\n",
    "        json_data.update({\"tool_choice\": function_call})\n",
    "\n",
    "    try:\n",
    "        response = requests.post(\"https://vip.apiyi.com/v1/chat/completions\", headers = headers, json = json_data)\n",
    "        return response\n",
    "    except Exception as e:\n",
    "        print(f\"Exception:{e}\")\n",
    "        return e\n",
    "\n",
    "# 使用colored打印消息对话内容\n",
    "def prety_print_vonversation(massages):\n",
    "    # 为不同角色设置不同颜色\n",
    "    role_to_color = {\n",
    "        \"developer\": \"red\",\n",
    "        \"user\": \"green\",\n",
    "        \"assistant\": \"blue\",\n",
    "        \"tool\": \"magenta\"\n",
    "    }\n",
    "    # 遍历消息列表，并用不同颜色输出\n",
    "    for message in messages:\n",
    "        if message[\"role\"] == \"developer\":\n",
    "            print(colored(f\"developer：{message['content']}\\n\", role_to_color[message[\"role\"]]))\n",
    "        elif message[\"role\"] == \"user\":\n",
    "            print(colored(f\"user：{message['content']}\\n\", role_to_color[message[\"role\"]]))\n",
    "        elif message[\"role\"] == \"assistant\" and message.get(\"tool_calls\"):\n",
    "            print(colored(f\"assistant[function_call]：{message['tool_calls']}\\n\", role_to_color[message[\"role\"]]))\n",
    "        elif message[\"role\"] == \"assistant\" and not message.get(\"tool_calls\"):\n",
    "            print(colored(f\"assistant[content]：{message['content']}\\n\", role_to_color[message[\"role\"]]))\n",
    "        elif message[\"role\"] == \"tool\":\n",
    "            print(colored(f\"function:{message}\\n\", role_to_color[message[\"role\"]]))\n",
    "\n",
    "# 拼装message消息体\n",
    "def splice_message(role, content):\n",
    "    return {\"role\": role, \"content\": content}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "7532dd01-1431-4794-a71d-aa5072b8dfc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义functions参数，包含两个方法的说明字典\n",
    "tools = [\n",
    "  {\n",
    "    \"type\": \"function\",\n",
    "    \"function\": {\n",
    "      \"name\": \"get_current_weather\",\n",
    "      \"description\": \"Get the current weather in a given location\",\n",
    "      \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "          \"location\": {\n",
    "            \"type\": \"string\",\n",
    "            \"description\": \"The city and state, e.g. San Francisco, CA\",\n",
    "          },\n",
    "          \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The temperature unit to use. Infer this from the users location.\"},\n",
    "        },\n",
    "        \"required\": [\"location\", \"unit\"],\n",
    "      },\n",
    "    }\n",
    "  },\n",
    "    {\n",
    "    \"type\": \"function\",\n",
    "    \"function\": {\n",
    "      \"name\": \"get_n_day_weather_forecast\",\n",
    "      \"description\": \"Get an N-day weather in a given location\",\n",
    "      \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "          \"location\": {\n",
    "            \"type\": \"string\",\n",
    "            \"description\": \"The city and state, e.g. San Francisco, CA\",\n",
    "          },\n",
    "          \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The temperature unit to use. Infer this from the users location.\"},\n",
    "          \"num_days\": {\"type\": \"integer\", \"description\": \"The number of days to forecast\"}\n",
    "        },\n",
    "        \"required\": [\"location\", \"unit\", \"num_days\"]\n",
    "      },\n",
    "    }\n",
    "  }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "abdee6e2-1f5d-43fe-a7f5-d6061de491c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mdeveloper：Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\n",
      "\u001b[0m\n",
      "\u001b[32muser：今天的天气怎么样？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[content]：请提供您所在的城市和州（或省）名称，以便我可以为您查询今天的天气。\n",
      "\u001b[0m\n",
      "\u001b[32muser：我现在在中国上海。\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_GsbO3qZPijFQVJjmTq7iN0c7', 'type': 'function', 'function': {'name': 'get_current_weather', 'parameters': None, 'arguments': '{\"location\":\"Shanghai, China\",\"unit\":\"celsius\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_GsbO3qZPijFQVJjmTq7iN0c7', 'name': 'get_current_weather', 'content': '{\"temperature\": \"33°C\", \"condition\": \"Cloudy\"}'}\n",
      "\u001b[0m\n",
      "\u001b[32muser：上海未来5天的天气会有什么变化吗？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_bAE0avyFs3okQPizk9u4Sjij', 'type': 'function', 'function': {'name': 'get_n_day_weather_forecast', 'parameters': None, 'arguments': '{\"location\":\"Shanghai, China\",\"num_days\":5,\"unit\":\"celsius\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_bAE0avyFs3okQPizk9u4Sjij', 'name': 'get_n_day_weather_forecast', 'content': '{\"temperature\": \"11°C\", \"condition\": \"Cloudy\"}'}\n",
      "\u001b[0m\n",
      "\u001b[32muser：我过两天要去北京，北京天气如何？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[content]：请问您去北京的具体日期是什么时候？这样我可以为您查询那天的天气情况。\n",
      "\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# 循环询问GPT，以获得答案\n",
    "# messages = []\n",
    "\n",
    "# messages.append(splice_message(\"developer\", \"Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\"))\n",
    "\n",
    "# messages.append(splice_message(\"user\", \"今天的天气怎么样？\"))\n",
    "# messages.pop()\n",
    "# messages.append(splice_message(\"user\", \"我现在在中国上海。\"))\n",
    "\n",
    "# messages.append(splice_message(\"user\", \"我希望是摄氏度，在我没有特殊说明的时候，均为摄氏度\"))\n",
    "# messages.append(splice_message(\"user\", \"上海未来5天的天气会有什么变化吗？\"))\n",
    "messages.append(splice_message(\"user\", \"我过两天要去北京，北京天气如何？\"))\n",
    "\n",
    "# 调用接口，获取GPT生成结果\n",
    "chat_response = chat_completion_request(messages, functions=tools)\n",
    "assistant_message = chat_response.json()[\"choices\"][0][\"message\"]\n",
    "# print(assistant_message)\n",
    "\n",
    "# 将回答作为下一次的输入上下文\n",
    "messages.append(assistant_message)\n",
    "\n",
    "# 如果是当前问题已回答完毕，function_call有返回参数，如果需要继续借助该上下文，就必须设置tool的补全结果，否则会出现异常\n",
    "if chat_response.json()[\"choices\"][0][\"finish_reason\"] in (\"tool_calls\", \"stop\"):\n",
    "    messages.append({\"role\": \"tool\", \"tool_call_id\": assistant_message['tool_calls'][0][\"id\"], \"name\": assistant_message['tool_calls'][0][\"function\"][\"name\"], \"content\": \"{\\\"temperature\\\": \\\"\" + str(random.randint(-10,38)) + \"°C\\\", \\\"condition\\\": \\\"Cloudy\\\"}\"})\n",
    "# 输出消息记录\n",
    "prety_print_vonversation(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "13aa977d-0de6-483e-9e67-83bf5a8c5997",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mdeveloper：Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\n",
      "\u001b[0m\n",
      "\u001b[32muser：今天的天气怎么样？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[content]：请提供您所在的城市和州（或省）名称，以便我可以为您查询今天的天气。\n",
      "\u001b[0m\n",
      "\u001b[32muser：我现在在中国上海。\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_GsbO3qZPijFQVJjmTq7iN0c7', 'type': 'function', 'function': {'name': 'get_current_weather', 'parameters': None, 'arguments': '{\"location\":\"Shanghai, China\",\"unit\":\"celsius\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_GsbO3qZPijFQVJjmTq7iN0c7', 'name': 'get_current_weather', 'content': '{\"temperature\": \"33°C\", \"condition\": \"Cloudy\"}'}\n",
      "\u001b[0m\n",
      "\u001b[32muser：上海未来5天的天气会有什么变化吗？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_bAE0avyFs3okQPizk9u4Sjij', 'type': 'function', 'function': {'name': 'get_n_day_weather_forecast', 'parameters': None, 'arguments': '{\"location\":\"Shanghai, China\",\"num_days\":5,\"unit\":\"celsius\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_bAE0avyFs3okQPizk9u4Sjij', 'name': 'get_n_day_weather_forecast', 'content': '{\"temperature\": \"11°C\", \"condition\": \"Cloudy\"}'}\n",
      "\u001b[0m\n",
      "\u001b[32muser：我过两天要去北京，北京天气如何？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_96sowD4C6Cu8Wv6n9Pg7WA2M', 'type': 'function', 'function': {'name': 'get_n_day_weather_forecast', 'parameters': None, 'arguments': '{\"location\":\"Beijing, China\",\"num_days\":2,\"unit\":\"celsius\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_96sowD4C6Cu8Wv6n9Pg7WA2M', 'name': 'get_n_day_weather_forecast', 'content': '{\"temperature\": \"10°C\", \"condition\": \"Cloudy\"}'}\n",
      "\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# messages.append(splice_message(\"user\", \"我过两天要去北京，北京天气如何？\"))\n",
    "\n",
    "chat_response = chat_completion_request(messages, functions=tools, function_call={\"type\": \"function\", \"function\": {\"name\": \"get_n_day_weather_forecast\"}})\n",
    "\n",
    "assistant_message = chat_response.json()[\"choices\"][0][\"message\"]\n",
    "messages.append(assistant_message)\n",
    "\n",
    "if chat_response.json()[\"choices\"][0][\"finish_reason\"] in (\"tool_calls\", \"stop\"):\n",
    "    messages.append({\"role\": \"tool\", \"tool_call_id\": assistant_message['tool_calls'][0][\"id\"], \"name\": assistant_message['tool_calls'][0][\"function\"][\"name\"], \"content\": \"{\\\"temperature\\\": \\\"\" + str(random.randint(-10,38)) + \"°C\\\", \\\"condition\\\": \\\"Cloudy\\\"}\"})\n",
    "\n",
    "prety_print_vonversation(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "5045ae67-a6e0-4bf4-9c19-a74ff9461942",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mdeveloper：Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\n",
      "\u001b[0m\n",
      "\u001b[32muser：今天的天气怎么样？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[content]：请提供您所在的城市和州（或省）名称，以便我可以为您查询今天的天气。\n",
      "\u001b[0m\n",
      "\u001b[32muser：我现在在中国上海。\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_GsbO3qZPijFQVJjmTq7iN0c7', 'type': 'function', 'function': {'name': 'get_current_weather', 'parameters': None, 'arguments': '{\"location\":\"Shanghai, China\",\"unit\":\"celsius\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_GsbO3qZPijFQVJjmTq7iN0c7', 'name': 'get_current_weather', 'content': '{\"temperature\": \"33°C\", \"condition\": \"Cloudy\"}'}\n",
      "\u001b[0m\n",
      "\u001b[32muser：上海未来5天的天气会有什么变化吗？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_bAE0avyFs3okQPizk9u4Sjij', 'type': 'function', 'function': {'name': 'get_n_day_weather_forecast', 'parameters': None, 'arguments': '{\"location\":\"Shanghai, China\",\"num_days\":5,\"unit\":\"celsius\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_bAE0avyFs3okQPizk9u4Sjij', 'name': 'get_n_day_weather_forecast', 'content': '{\"temperature\": \"11°C\", \"condition\": \"Cloudy\"}'}\n",
      "\u001b[0m\n",
      "\u001b[32muser：我过两天要去北京，北京天气如何？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_96sowD4C6Cu8Wv6n9Pg7WA2M', 'type': 'function', 'function': {'name': 'get_n_day_weather_forecast', 'parameters': None, 'arguments': '{\"location\":\"Beijing, China\",\"num_days\":2,\"unit\":\"celsius\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_96sowD4C6Cu8Wv6n9Pg7WA2M', 'name': 'get_n_day_weather_forecast', 'content': '{\"temperature\": \"10°C\", \"condition\": \"Cloudy\"}'}\n",
      "\u001b[0m\n",
      "\u001b[32muser：2天后，我要去北京，那边天气如何？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[content]：请稍等，我将为您查询北京两天后的天气。\n",
      "\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "messages.append(splice_message(\"user\", \"2天后，我要去北京，那边天气如何？\"))\n",
    "\n",
    "chat_response = chat_completion_request(messages, functions=tools, function_call=\"none\")\n",
    "\n",
    "assistant_message = chat_response.json()[\"choices\"][0][\"message\"]\n",
    "messages.append(assistant_message)\n",
    "\n",
    "if chat_response.json()[\"choices\"][0][\"finish_reason\"] in (\"tool_calls\", \"stop\") and assistant_message.get(\"tool_calls\"):\n",
    "    messages.append({\"role\": \"tool\", \"tool_call_id\": assistant_message['tool_calls'][0][\"id\"], \"name\": assistant_message['tool_calls'][0][\"function\"][\"name\"], \"content\": \"{\\\"temperature\\\": \\\"\" + str(random.randint(-10,38)) + \"°C\\\", \\\"condition\\\": \\\"Cloudy\\\"}\"})\n",
    "\n",
    "prety_print_vonversation(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67c20b1f-9200-4c67-986e-01694f59bffd",
   "metadata": {},
   "source": [
    "## SQL查询案例"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "f4a72120-3805-4f6e-a41c-9d2cdfc0ede9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opened database successfully\n"
     ]
    }
   ],
   "source": [
    "import sqlite3\n",
    "\n",
    "conn = sqlite3.connect(\"data/chinook.db\")\n",
    "print(\"Opened database successfully\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "9bab7c44-e00e-4ba9-b0b8-259c65795362",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_table_names(conn):\n",
    "    table_names = []\n",
    "    tables = conn.execute(\"select name from sqlite_master where type = 'table';\")\n",
    "    for table in tables.fetchall():\n",
    "        table_names.append(table[0])\n",
    "    return table_names\n",
    "\n",
    "def get_colunm_names(conn, table_name):\n",
    "    column_names = []\n",
    "    columns = conn.execute(f\"pragma table_info('{table_name}');\").fetchall()\n",
    "    for col in columns:\n",
    "        column_names.append(col[1])\n",
    "    return column_names\n",
    "\n",
    "def get_database_info(conn):\n",
    "    table_dicts = []\n",
    "    for table_name in get_table_names(conn):\n",
    "        column_names = get_colunm_names(conn, table_name)\n",
    "        table_dicts.append({\"table_name\": table_name, \"column_names\": column_names})\n",
    "    return table_dicts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "4161554d-c52f-4eff-baa7-33f3c3f454f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "database_schema_dict = get_database_info(conn)\n",
    "# database_schema_dict\n",
    "\n",
    "database_schema_string = \"\\n\".join([f\"Table: {table['table_name']}\\nColumns: {', '.join(table['column_names'])}\" for table in database_schema_dict])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "2fc14de1-7293-4a2e-9d22-b3b1e93de0bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'table_name': 'albums', 'column_names': ['AlbumId', 'Title', 'ArtistId']}, {'table_name': 'sqlite_sequence', 'column_names': ['name', 'seq']}, {'table_name': 'artists', 'column_names': ['ArtistId', 'Name']}, {'table_name': 'customers', 'column_names': ['CustomerId', 'FirstName', 'LastName', 'Company', 'Address', 'City', 'State', 'Country', 'PostalCode', 'Phone', 'Fax', 'Email', 'SupportRepId']}, {'table_name': 'employees', 'column_names': ['EmployeeId', 'LastName', 'FirstName', 'Title', 'ReportsTo', 'BirthDate', 'HireDate', 'Address', 'City', 'State', 'Country', 'PostalCode', 'Phone', 'Fax', 'Email']}, {'table_name': 'genres', 'column_names': ['GenreId', 'Name']}, {'table_name': 'invoices', 'column_names': ['InvoiceId', 'CustomerId', 'InvoiceDate', 'BillingAddress', 'BillingCity', 'BillingState', 'BillingCountry', 'BillingPostalCode', 'Total']}, {'table_name': 'invoice_items', 'column_names': ['InvoiceLineId', 'InvoiceId', 'TrackId', 'UnitPrice', 'Quantity']}, {'table_name': 'media_types', 'column_names': ['MediaTypeId', 'Name']}, {'table_name': 'playlists', 'column_names': ['PlaylistId', 'Name']}, {'table_name': 'playlist_track', 'column_names': ['PlaylistId', 'TrackId']}, {'table_name': 'tracks', 'column_names': ['TrackId', 'Name', 'AlbumId', 'MediaTypeId', 'GenreId', 'Composer', 'Milliseconds', 'Bytes', 'UnitPrice']}, {'table_name': 'sqlite_stat1', 'column_names': ['tbl', 'idx', 'stat']}]"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "database_schema_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "0ceb0a95-fce8-49bc-896f-2bdb1923cdb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "functions = [\n",
    "  {\n",
    "    \"type\": \"function\",\n",
    "    \"function\": {\n",
    "      \"name\": \"ask_database\",\n",
    "      \"description\": \"Use this function to answer user questions about music. Output should be a fully formed SQL query.\",\n",
    "      \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "          \"query\": {\n",
    "            \"type\": \"string\",\n",
    "            \"description\": f\"\"\"\n",
    "                            SQL query extracting info to answer the user's question.\n",
    "                            SQL should be written using this database schema:\n",
    "                            {database_schema_string}\n",
    "                            The query should be returned in plain text, not in JSON.\n",
    "                            \"\"\"\n",
    "          }\n",
    "        },\n",
    "        \"required\": [\"query\"],\n",
    "      },\n",
    "    }\n",
    "  }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "6268318d-700b-4f55-8f8c-9ff70103a7b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask_database(conn, query):\n",
    "    try:\n",
    "        results = str(conn.execute(query).fetchall())\n",
    "    except Exception as e:\n",
    "        results = f\"query failed with error: {e}\"\n",
    "    return results\n",
    "\n",
    "def execute_function_call(massage):\n",
    "    if massage['tool_calls'][0][\"function\"][\"name\"] == \"ask_database\":\n",
    "        query = json.loads(massage['tool_calls'][0][\"function\"][\"arguments\"])[\"query\"]\n",
    "        results = ask_database(conn, query)\n",
    "    else:\n",
    "        results = f\"Error: function {massage['tool_calls'][0]['function']['name']} does not exist\"\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "c880c30f-734d-44d3-b29d-3c158dd97e97",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mdeveloper：Answer user questions by generating SQL queries against the Chinook Music Database.\n",
      "\u001b[0m\n",
      "\u001b[32muser：你好，谁是按曲目数量排名前5的艺术家？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'id': 'call_AnoXEKAvnWV8yOh2dNayCYuY', 'type': 'function', 'function': {'name': 'ask_database', 'parameters': None, 'arguments': '{\"query\":\"SELECT artists.Name, COUNT(tracks.TrackId) AS TrackCount\\\\nFROM artists\\\\nJOIN albums ON artists.ArtistId = albums.ArtistId\\\\nJOIN tracks ON albums.AlbumId = tracks.AlbumId\\\\nGROUP BY artists.ArtistId\\\\nORDER BY TrackCount DESC\\\\nLIMIT 5;\"}'}}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_AnoXEKAvnWV8yOh2dNayCYuY', 'name': 'ask_database', 'content': \"[('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Deep Purple', 92)]\"}\n",
      "\u001b[0m\n",
      "\u001b[32muser：曲目最多的专辑叫什么名字？\n",
      "\u001b[0m\n",
      "\u001b[34massistant[function_call]：[{'function': {'arguments': '{\"query\":\"SELECT albums.Title, COUNT(tracks.TrackId) AS TrackCount\\\\nFROM albums\\\\nJOIN tracks ON albums.AlbumId = tracks.AlbumId\\\\nGROUP BY albums.AlbumId\\\\nORDER BY TrackCount DESC\\\\nLIMIT 1;\"}', 'name': 'ask_database'}, 'id': 'call_8aAeX8oeKqwrr5nJKfpnRYNR', 'type': 'function'}]\n",
      "\u001b[0m\n",
      "\u001b[35mfunction:{'role': 'tool', 'tool_call_id': 'call_8aAeX8oeKqwrr5nJKfpnRYNR', 'name': 'ask_database', 'content': \"[('Greatest Hits', 57)]\"}\n",
      "\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# messages = []\n",
    "\n",
    "# messages.append(splice_message(\"developer\", \"Answer user questions by generating SQL queries against the Chinook Music Database.\"))\n",
    "\n",
    "# messages.append(splice_message(\"user\", \"你好，谁是按曲目数量排名前5的艺术家？\"))\n",
    "messages.append(splice_message(\"user\", \"曲目最多的专辑叫什么名字？\"))\n",
    "\n",
    "chat_response = chat_completion_request(messages, functions)\n",
    "\n",
    "assistant_message = chat_response.json()[\"choices\"][0][\"message\"]\n",
    "\n",
    "messages.append(assistant_message)\n",
    "\n",
    "if chat_response.json()[\"choices\"][0][\"finish_reason\"] in (\"tool_calls\", \"stop\") and assistant_message.get(\"tool_calls\"):\n",
    "    results = execute_function_call(assistant_message)\n",
    "    messages.append({\"role\": \"tool\", \"tool_call_id\": assistant_message['tool_calls'][0][\"id\"], \"name\": assistant_message['tool_calls'][0][\"function\"][\"name\"], \"content\": results})\n",
    "\n",
    "prety_print_vonversation(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "29bbc084-3833-41a2-9e2e-0fcd249540d9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Deep Purple', 92)]\""
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = execute_function_call(assistant_message)\n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "387707db-935e-4ae9-8875-373e7c6e12ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'role': 'assistant', 'tool_calls': [{'id': 'call_LOqj0W0weEoLpuJAx59VgSSv', 'type': 'function', 'function': {'name': 'ask_database', 'parameters': None, 'arguments': '{\"query\":\"SELECT a.Name, COUNT(t.TrackId) AS TrackCount \\\\nFROM artists a \\\\nJOIN albums al ON a.ArtistId = al.ArtistId \\\\nJOIN tracks t ON al.AlbumId = t.AlbumId \\\\nGROUP BY a.ArtistId \\\\nORDER BY TrackCount DESC \\\\nLIMIT 5;\"}'}}]}"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc52a32f-03fe-498c-a3d9-2500bd434bce",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langChain",
   "language": "python",
   "name": "langchain"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
